import re
import graphviz
import sys

class ASTNode:
    def __init__(self, name, label=None, unique_id=None):
        self.name = name  # 节点名称
        self.label = label if label else name  # 默认标签为节点名称
        self.unique_id = unique_id  # 用于区分相同标签的节点
        self.children = []  # 子节点列表

    def add_child(self, child):
        """ 添加子节点 """
        self.children.append(child)

    def to_dot(self, parent=None, seen_edges=set()):
        """ 将当前节点转为 dot 格式，返回节点字符串和边的字符串 """
        node_str = f'  "{self.unique_id}" [label="{self.label}"];\n'
        edges_str = ""

        # 如果父节点存在且不是当前节点本身，检查是否已经生成过这条边
        if parent and parent.unique_id != self.unique_id:
            edge = (parent.unique_id, self.unique_id)
            if edge not in seen_edges:
                edges_str = f'  "{parent.unique_id}" -> "{self.unique_id}";\n'
                seen_edges.add(edge)
        
        # 遍历所有子节点
        for child in self.children:
            child_str, child_edges = child.to_dot(parent=self, seen_edges=seen_edges)
            node_str += child_str
            edges_str += child_edges

        return node_str, edges_str


def parse_input(input_str):
    """ 从输入字符串中提取父子节点关系，并生成AST树 """
    lines = input_str.strip().split('\n')
    node_stack = []  
    unique_id_counter = 0  

    for line in lines:
        indentation_level = line.count('|')
        node_name = line.split('>--')[1].strip() 
        current_node = ASTNode(node_name, unique_id=unique_id_counter)

        while len(node_stack) > indentation_level:
            node_stack.pop()

        if node_stack:
            node_stack[-1].add_child(current_node)

        node_stack.append(current_node)

        unique_id_counter += 1

    return node_stack[0]


def generate_ast_from_file(input_file):
    """ 从文件读取输入并生成AST图形 """
    with open(input_file, 'r') as file:
        input_data = file.read()

    root_node = parse_input(input_data)

    graph = "digraph AST {\n"
    node_str, edges_str = root_node.to_dot()
    graph += node_str + edges_str
    graph += "}"

    dot_file = "ast_graph.dot"
    with open(dot_file, "w") as file:
        file.write(graph)

    print(f"DOT图已保存为 {dot_file} 文件。")

    with open(dot_file, "r") as file:
        dot_content = file.read()

    graph = graphviz.Source(dot_content)
    graph.format = 'png'
    graph.render('ast_graph', cleanup=True)

    print("PNG 图片已保存为 ast_graph.png 文件。")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("""
Usage: python ast_to_png.py <input_file>

Arguments:
  <input_file>  A text file containing the AST node structure. The file should follow a specific tree-like structure where nodes are connected with '>' and '|' characters.

""")
        sys.exit(1)

    input_file = sys.argv[1]
    generate_ast_from_file(input_file)